package rabbitmq

import (
	"data-center/config"
	"data-center/pkg/queue"
	"fmt"
	"github.com/streadway/amqp"
	"time"
)

func init() {
	// 注册工厂函数
	queue.RegisterRabbitMQProducer(NewProducer)
	queue.RegisterRabbitMQConsumer(NewConsumer)
}

// Producer RabbitMQ生产者
type Producer struct {
	conn     *amqp.Connection
	channel  *amqp.Channel
	exchange string
	queues   map[string]string
}

// Consumer RabbitMQ消费者
type Consumer struct {
	conn     *amqp.Connection
	channel  *amqp.Channel
	delivery <-chan amqp.Delivery
}

// NewProducer 创建RabbitMQ生产者
func NewProducer(cfg config.QueueConfig) (queue.Producer, error) {
	conn, err := amqp.Dial(cfg.RabbitMQ.URI)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to RabbitMQ: %w", err)
	}

	ch, err := conn.Channel()
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to open a channel: %w", err)
	}

	// 声明交换机
	err = ch.ExchangeDeclare(
		cfg.RabbitMQ.Exchange, // name
		"direct",              // type
		true,                  // durable
		false,                 // auto-deleted
		false,                 // internal
		false,                 // no-wait
		nil,                   // arguments
	)
	if err != nil {
		ch.Close()
		conn.Close()
		return nil, fmt.Errorf("failed to declare an exchange: %w", err)
	}

	// 声明队列
	queues := map[string]string{
		"task":  cfg.Queues.Task,
		"file":  cfg.Queues.File,
		"index": cfg.Queues.Index,
	}

	for _, queueName := range queues {
		_, err = ch.QueueDeclare(
			queueName, // name
			true,      // durable
			false,     // delete when unused
			false,     // exclusive
			false,     // no-wait
			nil,       // arguments
		)
		if err != nil {
			ch.Close()
			conn.Close()
			return nil, fmt.Errorf("failed to declare a queue: %w", err)
		}

		// 绑定队列到交换机
		err = ch.QueueBind(
			queueName,             // queue name
			queueName,             // routing key
			cfg.RabbitMQ.Exchange, // exchange
			false,
			nil,
		)
		if err != nil {
			ch.Close()
			conn.Close()
			return nil, fmt.Errorf("failed to bind a queue: %w", err)
		}
	}

	return &Producer{
		conn:     conn,
		channel:  ch,
		exchange: cfg.RabbitMQ.Exchange,
		queues:   queues,
	}, nil
}

// NewConsumer 创建RabbitMQ消费者
func NewConsumer(cfg config.QueueConfig, queueTypes []string) (queue.Consumer, error) {
	conn, err := amqp.Dial(cfg.RabbitMQ.URI)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to RabbitMQ: %w", err)
	}

	ch, err := conn.Channel()
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to open a channel: %w", err)
	}

	// 确定要消费的队列
	var queueName string
	if len(queueTypes) > 0 {
		switch queueTypes[0] {
		case "task":
			queueName = cfg.Queues.Task
		case "file":
			queueName = cfg.Queues.File
		case "index":
			queueName = cfg.Queues.Index
		default:
			ch.Close()
			conn.Close()
			return nil, fmt.Errorf("unknown queue type: %s", queueTypes[0])
		}
	} else {
		ch.Close()
		conn.Close()
		return nil, fmt.Errorf("no queue type specified")
	}

	// 声明队列
	_, err = ch.QueueDeclare(
		queueName, // name
		true,      // durable
		false,     // delete when unused
		false,     // exclusive
		false,     // no-wait
		nil,       // arguments
	)
	if err != nil {
		ch.Close()
		conn.Close()
		return nil, fmt.Errorf("failed to declare a queue: %w", err)
	}

	// 消费消息
	delivery, err := ch.Consume(
		queueName, // queue
		"",        // consumer
		false,     // auto-ack
		false,     // exclusive
		false,     // no-local
		false,     // no-wait
		nil,       // args
	)
	if err != nil {
		ch.Close()
		conn.Close()
		return nil, fmt.Errorf("failed to register a consumer: %w", err)
	}

	return &Consumer{
		conn:     conn,
		channel:  ch,
		delivery: delivery,
	}, nil
}

// Publish 发布消息到指定队列
func (p *Producer) Publish(queueType string, message []byte) error {
	queueName, ok := p.queues[queueType]
	if !ok {
		return fmt.Errorf("unknown queue type: %s", queueType)
	}

	return p.channel.Publish(
		p.exchange, // exchange
		queueName,  // routing key
		false,      // mandatory
		false,      // immediate
		amqp.Publishing{
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	)
}

// Consume 消费消息
func (c *Consumer) Consume(timeout int) ([]byte, error) {
	select {
	case msg := <-c.delivery:
		msg.Ack(false)
		return msg.Body, nil
	case <-time.After(time.Duration(timeout) * time.Millisecond):
		return nil, fmt.Errorf("timeout waiting for message")
	}
}

// Close 关闭连接
func (p *Producer) Close() {
	p.channel.Close()
	p.conn.Close()
}

// Close 关闭连接
func (c *Consumer) Close() {
	c.channel.Close()
	c.conn.Close()
}
